#include<vector>
#include<iostream>
#include<fstream>
#include<functional>
#include<cmath>
using namespace std;

function<double(double)> actF;

class ne{
	vector<ne*> fa;
	vector<double> w;
	bool ed;
	double rs;
public:
	double delta;
	ne(){
		ed=0;
		delta=0;
	}
	virtual double count(){
		if(ed){
			return rs;
		}
		rs=0;
		ed=1;
		for(auto &a:fa){
			rs+=a->count()*w[(&a-&fa[0])];
		}
		rs=actF(rs);
		return rs;
	}
	void connect(ne &fn,double fw){
		fa.push_back(&fn);
		w.push_back(fw);
	}
	void reset(){
		ed=0;
		delta=0;
	}
	friend class nenet;
};

class input:public ne{
public:
	double va;
	double count(){
		return actF(va);
	}
};

class nenet{
	input in[200];
	ne n[19800];
	const double lr=0.1;//学习率
public:
	string md;
	static void Ninit(){
		actF=[](double d){
			return (d>0)?d:0;//d/(1+exp(-d));
		};
	}

	// 优化：二进制加载模型（移除文本转换，按顺序读权重）
	void init(){
		// 二进制模式打开文件，ios::binary不可或缺
		ifstream ifp(md, ios::in | ios::binary);
		if(!ifp.is_open()){
			cerr<<"[Error] Failed to open model file: "<<md<<endl;
			exit(1);
		}

		// 二进制读取封装（直接操作内存，无转换开销）
		auto inf=[&]{
			double dv;
			ifp.read(reinterpret_cast<char*>(&dv), sizeof(double));
			return dv;
		};

		// 1. 加载输入层→第一层权重（顺序与保存时一致）
		for(int i=0;i<200;i++){
			for(int j=0;j<200;j++){
				n[i].connect(in[j], inf());
			}
		}

		// 2. 加载隐藏层之间权重（按连接顺序直接读，无需找索引）
		for(int k=0;k<98;k++){
			int currLayerStart = k*200;
			int nextLayerStart = (k+1)*200;
			for(int m=nextLayerStart;m<nextLayerStart+200;m++){
				// 连接时p的顺序 = w数组的顺序，直接按顺序读
				for(int p=currLayerStart;p<currLayerStart+200;p++){
					n[m].connect(n[p], inf());
				}
			}
		}
		ifp.close();
	}

	// 优化：二进制保存模型（移除索引查找，按顺序写权重）
	void store(){
		// 二进制模式写入，ios::binary不可或缺
		ofstream ofp(md, ios::out | ios::binary | ios::trunc);
		if(!ofp.is_open()){
			cerr<<"[Error] Failed to write model file: "<<md<<endl;
			exit(1);
		}

		// 二进制写入封装（直接写内存数据）
		auto outf=[&](double dv){
			ofp.write(reinterpret_cast<const char*>(&dv), sizeof(double));
		};

		// 1. 保存输入层→第一层权重（w顺序与连接顺序一致）
		for(int i=0;i<200;i++){
			for(int j=0;j<200;j++){
				outf(n[i].w[j]);
			}
		}

		// 2. 保存隐藏层之间权重（直接按w顺序写，无需找索引）
		for(int k=0;k<98;k++){
			int currLayerStart = k*200;
			int nextLayerStart = (k+1)*200;
			for(int m=nextLayerStart;m<nextLayerStart+200;m++){
				// w数组顺序 = 连接时p的顺序，直接遍历w即可
				for(auto &weight : n[m].w){
					outf(weight);
				}
			}
		}
		ofp.close();
	}

	void init1st(){
		for(int i=0;i<200;i++){
			for(int j=0;j<200;j++){
				n[i].connect(in[j],(rand()%1000)/1000.0-0.5);
			}
		}
		for(int k=0;k<98;k++){
			int currLayerStart = k*200;
			int nextLayerStart = (k+1)*200;
			for(int m=nextLayerStart;m<nextLayerStart+200;m++){
				for(int p=currLayerStart;p<currLayerStart+200;p++){
					n[m].connect(n[p],(rand()%1000)/1000.0-0.5);
				}
			}
		}
	}

	string ask(string qu){
		for(auto &nn:n) nn.reset();
		for(int i=0;i<200;i++){
			char c=qu[i%qu.size()];
			in[i].va=(unsigned char)c/255.0;
		}
		string res;
		for(int i=19600;i<19800;i++){
			double o=n[i].count();
			char ch=(char)(o*255.0);
			//if(!ch) break;
			res+=ch;
		}
		return res;
	}

// 只修改train函数，其他代码保持不变
	void train(string question, string expect){
		string out = ask(question);
		vector<double> target(200, 0);
		for(int i = 0; i < 200; ++i) {
			char c = expect[i % expect.size()];
			target[i] = (unsigned char)c / 255.0;
		}
	
		// 前向计算时缓存所有输出，避免重复计算
		vector<double> outputs(19800);
		for(int i = 0; i < 19800; ++i) {
			outputs[i] = n[i].count();
		}
	
		// 输出层误差计算 - 优化：使用缓存的输出值
		int outputLayerStart = 19600;
		for(int i = outputLayerStart; i < 19800; ++i) {
			int idx = i - outputLayerStart;
			double o = outputs[i];
			// 针对ReLU激活函数优化导数计算
			double deriv = (o > 0) ? 1.0 : 0.0;
			n[i].delta = (target[idx] - o) * deriv;
		}
	
		// 隐藏层误差计算 - 优化：预计算下层索引，消除查找
		for(int k = 97; k >= 0; --k) {
			int currLayerStart = k * 200;
			int nextLayerStart = (k + 1) * 200;
			
			// 预计算当前层每个神经元的误差贡献
			vector<double> layerDeltas(200, 0.0);
			
			// 遍历下一层神经元，累积误差
			for(int j = nextLayerStart; j < nextLayerStart + 200; ++j) {
				// 直接通过索引访问权重，避免指针比较
				for(int p = 0; p < 200; ++p) {
					int currIdx = p; // 当前神经元在前一层中的索引
					layerDeltas[currIdx] += n[j].delta * n[j].w[p];
				}
			}
			
			// 计算当前层每个神经元的delta
			for(int i = 0; i < 200; ++i) {
				int neuronIdx = currLayerStart + i;
				double o = outputs[neuronIdx];
				double deriv = (o > 0) ? 1.0 : 0.0;
				n[neuronIdx].delta = layerDeltas[i] * deriv;
			}
		}
	
		// 权重更新 - 优化：减少函数调用和指针操作
		for(int k = 98; k >= 0; --k) {
			int currLayerStart = k * 200;
			for(int i = currLayerStart; i < currLayerStart + 200; ++i) {
				// 直接使用预计算的输出值，避免重复调用count()
				for(int j = 0; j < n[i].fa.size(); ++j) {
					ne* prev = n[i].fa[j];
					double prevOutput;
					
					// 区分输入层和隐藏层，输入层有va属性
					if(prev >= in && prev < in + 200) {
						prevOutput = actF(static_cast<input*>(prev)->va);
					} else {
						int prevIdx = prev - n; // 计算索引
						prevOutput = outputs[prevIdx];
					}
					
					n[i].w[j] += lr * n[i].delta * prevOutput;
				}
			}
		}
	}

};

int main(int ac,char **av){
	auto help=[]{
		cout<<R"(Welcome to use jett-ai!
Jett-ai is an AI written in C++.

Usage: jett-ai [options] [question]

Options:
  -h, --help           Show this help message
  -q, --question <str> Specify the question to ask the AI(can be omitted)
  -t, --train <num>    Set training iterations (default: 0)
  -e, --expect <str>   Specify expected output for training
  -m, --model <model>  Set jett-ai model to use (default: std)
  -n, --new <model>    New a jett-ai model
Examples:
  jett -q "Hello, what is your name?"
  jett "Hey, how is every thing going?"
  jett -t 1000 -e "positive" "Is this good?" -m std
)";
		return 0;
	};
	if(ac==1){
		return help();
	}
	nenet::Ninit();
	
	string qu="[none]",ex,mo="std";
	int tr=0;bool nw=0;
	
	for(int i=1;i<ac;i++){
		if(av[i][0]=='-'){
			string avi=av[i];
			if(avi=="-q"||avi=="--question"){
				qu=av[++i];
			}else if(avi=="-t"||avi=="--train"){
				tr=atoi(av[++i]);
			}else if(avi=="-e"||avi=="--expect"){
				ex=av[++i];
			}else if(avi=="-h"||avi=="--help"){
				help();
			}else if(avi=="-m"||avi=="--model"){
				mo=av[++i];
			}else if(avi=="-n"||avi=="--new"){
				mo=av[++i];
				nw=1;
			}
		}else{
			qu=av[i];
		}
	}
	nenet net;
	net.md=mo;
	
	// 新建模型逻辑（无改动）
	if(nw){
		ifstream ifp(mo);
		if(ifp.is_open()){
			cout<<"[Warning] "<<mo<<" already exists!\n";
			return 1;
		}
		net.init1st();
		net.store();// 新建时直接保存二进制模型
		cout<<"[Success] New model created: "<<mo<<endl;
		return 0;
	}else{
		// 加载二进制模型
		net.init();
	}
	
	// 训练/推理逻辑（无改动）
	if(qu=="[none]"){
		return 0;
	}
	if(tr){
		cout<<"[Training] Iterations: "<<tr<<" ...\n";
		int k,j;
		for(int i=1;i<=tr;i++){
			k=i*100/tr;
			j=k;
			while(k--){
				cout<<"-";
			}
			j=100-j;
			while(j--){
				cout<<" ";
			}
			cout<<"|"<<i<<"/"<<tr<<"\r";
			net.train(qu,ex);
		}
		net.store();
		cout<<"[Success] Model saved to: "<<mo<<endl;
	}else{
		cout<<net.ask(qu)<<"\n";
	}
	return 0;
}